#include "Types.hpp"


__kernel void setCellsToMinus(__global int* cells,
	      		      const uint NCells)
{ 
    uint i = get_global_id(0);
    if (i >= NCells) return;
    
    cells[i]=-1;	
}

__kernel void sortParticleIDs(__global int* particle_ids,
	      		      const uint N)
{
    uint i = get_global_id(0);
    if (i >= N) return;
    
    particle_ids[i]=i;
}

__kernel void assembleCells(const __global real4* positions,
	      		      __global int* part_ids,
			      __global int* cells,
			      const uint N,
			      const uint nx,
			      const uint ny,
			      const uint nz,
			      const real len_x,
			      const real len_y,
			      const real len_z)
{
    uint i = get_global_id(0);
    if (i >= N) return;
    
    int cx = positions[i].x/len_x;
    int cy = positions[i].y/len_y;
    int cz = positions[i].z/len_z;

    int cellID=cx*ny*nz+cy*nz+cz;
    part_ids[i]=atom_xchg(cells+cellID,i) ;
}

__kernel void forceCalculation(const __global real4* positions,
                               __global real4* forces,
                               __local real4* forceSums,
			       const __global int* cells,
			       const __global int* part_ids, 
			       const __global real* radius,
			       const __global real4* velocities,
                               const uint N,
			       const uint nx,
			       const uint ny,
			       const uint nz,
			       const real len_x,
			       const real len_y,
			       const real len_z,
			       const real ks,
			       const real kd,
			       const int reflectx,
			       const int reflecty,
			       const int reflectz,
			       const real xmin,
			       const real ymin,
			       const real zmin,
			       const real xmax,
			       const real ymax,
			       const real zmax, 
				   const __global real4* forcGrav )
{
  uint i = get_global_id(0);
  if (i >= N) return;
  
  int cx = positions[i].x/len_x;
  int cy = positions[i].y/len_y;
  int cz = positions[i].z/len_z;

  int cellID1=cx*ny*nz+cy*nz+cz;

  
  forceSums[get_local_id(0)] =(real4)0.0;
  
 for(int l=-1;l<=1;++l)
  {
	for(int m=-1;m<=1;++m)
	{
		for(int n=-1;n<=1;++n)
		{
			int neighbor_cx=cx+l;
			int neighbor_cy=cy+m;
			int neighbor_cz=cz+n;
			
			if(reflectx==0)
			{
				if(neighbor_cx<0)
				     neighbor_cx=nx-1;
				else if(neighbor_cx>nx-1)
			     	     neighbor_cx=0;
			}
			else
			{
				if(neighbor_cx<0 || neighbor_cx>nx-1)
					continue;		
			}	
			
			if(reflecty==0)
			{
				if(neighbor_cy<0)
				     neighbor_cy=ny-1;
			     	else if(neighbor_cy>ny-1)
			     	     neighbor_cy=0;
			}
			else
			{
				if(neighbor_cy<0 || neighbor_cy>ny-1)
					continue;   			     	
			}			
			
			if(reflectz==0)
			{
				if(neighbor_cz<0)
				     neighbor_cz=nz-1;
				else if(neighbor_cz>nz-1)
			    	     neighbor_cz=0;			
			}
			else
			{
				if(neighbor_cz<0 || neighbor_cz>nz-1)
			        	continue;
			}
			int cellID=neighbor_cx*ny*nz+neighbor_cy*nz+neighbor_cz;
			int j=cells[cellID];
			while(j!=-1)
			{
				real4 dij = positions[i] - positions[j];
      				dij.w=0.0;
				real n2sqr = dij.x*dij.x + dij.y*dij.y + dij.z*dij.z;
				real norm = sqrt(n2sqr);
				real p = radius[j]+radius[i]-norm;
				if(p>0)
				{
					dij = dij/norm;
					real4 dvij = velocities[j]-velocities[i];
					dvij.w=0.0;
					real pdot = dij.x*dvij.x + dij.y*dvij.y + dij.z*dvij.z;
					real Fn = ks*p + kd*pdot;
      					dij *= Fn;

    					if (i != j)
       				   	   forceSums[get_local_id(0)] += dij;
				}
				j=part_ids[j];
			}
		}
	}
  }
  real pw=0, pdotw=0, Fnw=0;
  real4 normalV;
  normalV.x=0.0;   normalV.y=0.0;   normalV.z=0.0;   normalV.w=0.0; 
  if(reflectx!=0)
  {
	normalV.x=1.0;
	normalV.y=0.0;
	normalV.z=0.0;
	pw=xmin-(positions[i].x-radius[i]);
	if(pw>0)
	{
		pdotw=-velocities[i].x;
		real Fnw = ks*pw + kd*pdotw;
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;
	}
	normalV.x=-1.0;
	normalV.y=0.0;
	normalV.z=0.0;
	pw = positions[i].x-xmax+radius[i];	
	if(pw>0)
	{
		pdotw=-velocities[i].x;
		Fnw = ks*pw + kd*pdotw;
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;
	}										
  }
  if(reflecty!=0)
  {
	normalV.x=0.0;
	normalV.y=1.0;
	normalV.z=0.0;
	pw=ymin-(positions[i].y-radius[i]);
	if(pw>0)
	{
		pdotw=-velocities[i].y;
		real Fnw = ks*pw + kd*pdotw;
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;
	}
	normalV.x=0.0;
	normalV.y=-1.0;
	normalV.z=0.0;
	pw = positions[i].y-ymax+radius[i];	
	if(pw>0)
	{
		pdotw=-velocities[i].y;
		Fnw = ks*pw + kd*pdotw;
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;	
	}											
  }
  if(reflectz!=0)
  {
	normalV.x=0.0;
	normalV.y=0.0;
	normalV.z=1.0;
	pw=zmin-(positions[i].z-radius[i]);
	if(pw>0)
	{
		pdotw=-velocities[i].z;
		real Fnw = ks*pw + kd*pdotw;
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;
	}
	normalV.x=0.0;
	normalV.y=0.0;
	normalV.z=-1.0;
	pw = positions[i].z-zmax+radius[i];	
	if(pw>0)
	{
		pdotw=-velocities[i].z;
		Fnw = ks*pw + kd*pdotw;	
		normalV *= Fnw;
		forceSums[get_local_id(0)]+= normalV;
	}										
  }
  forceSums[get_local_id(0)]+= forcGrav[i];
	
  forces[get_global_id(0)] = forceSums[get_local_id(0)];
}

__kernel void setPeriodicity( __global real4* positions,
	      		      const uint N,
			      const real xmin,
			      const real xmax,
			      const real ymin,
			      const real ymax,
			      const real zmin,
			      const real zmax,
				  const int reflectx,
			       const int reflecty,
			       const int reflectz)
{
    uint i = get_global_id(0);
    if (i >= N) return;
    
	if (reflectx==0)
	{
		if(positions[i].x<xmin)
		positions[i].x+=(xmax-xmin);
		else if(positions[i].x>xmax)
		positions[i].x-=(xmax-xmin);
	}

	if (reflecty==0)
	{
    
		if(positions[i].y<ymin)
		positions[i].y+=(ymax-ymin);
		else if(positions[i].y>ymax)
		positions[i].y-=(ymax-ymin);
	}

	if (reflectz==0)
	{
		if(positions[i].z<zmin)
		positions[i].z+=(zmax-zmin);
		else if(positions[i].z>zmax)
		positions[i].z-=(zmax-zmin);
	}
}


__kernel void updatePositions( __global real4* positions,
	      		       __global real4* old_forces,
			      const __global real4* forces,
                              const __global real4* velocities,
                              const __global real* inv_2masses,
                              const real dt,
                              const uint N)
{
    uint i = get_global_id(0);
    if (i >= N) return;

    positions[i]=positions[i]+dt*velocities[i]+forces[i]*dt*dt*inv_2masses[i];
    old_forces[i]=forces[i];
}

__kernel void updateVelocities(__global real4* velocities,
	      	               const __global real4* old_forces,
			       const __global real4* forces,
			       const __global real* inv_2masses,
                               const real dt,
                               const uint N)
{
    uint i = get_global_id(0);
    if (i >= N) return;
    velocities[i]=velocities[i]+(old_forces[i]+forces[i])*inv_2masses[i]*dt;
}
